package huffman;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * @author 谭继好
 * @Description
 * @date 2024-05-15 16:37
 */
public class HuffmanTree {

    // 原始字符串
    public char[] origin;
    // 哈夫曼树的根结点
    public HufNode root;
    // 编码对应关系
    public Map<Character, StringBuilder> relation = new HashMap<>();

    public Map<Character, StringBuilder> getRelation() {
        return relation;
    }

    // 根据字符串初始化哈夫曼树
    public HuffmanTree(String s) {
        Map<Character, Integer> map = new HashMap<>();
        char[] arr = s.toCharArray();
        this.origin = arr;

        // 统计字符出现次数
        for (int i = 0; i < arr.length; i++) {
            map.put(arr[i], map.getOrDefault(arr[i], 0) + 1);
        }

        if (map.isEmpty()) {
            root = null;
            // 所有数据放在叶子结点，即使只有一个元素也不能放在根结点，这里放在左子树
        } else if (map.size() == 1) {
            root = new HufNode(' ');
            char data = ' ';
            for (Character c : map.keySet()) {
                data = c;
            }
            root.left = new HufNode(data);
        } else {
            PriorityQueue<TempNode> queue = new PriorityQueue<>(new Comparator<TempNode>() {
                @Override
                public int compare(TempNode o1, TempNode o2) {
                    return o1.count - o2.count;
                }
            });
            // 从最小值开始取两个最小值依次构成左右子结点，构造父结点并添加到哈希表中
            // 便于后续该结点作为子结点。删除构造好的两个结点，并将新的数据插入到链表中
            char data = ' ';
            for (Character c : map.keySet()) {
                data = c;
                HufNode node = new HufNode(c);
                queue.add(new TempNode(map.get(c), node));
            }
            while (queue.size() > 1) {
                TempNode left = queue.poll(), right = queue.poll();
                queue.add(new TempNode(left.count + right.count, new HufNode(' ', left.node, right.node)));
            }
            root = queue.poll().node;
        }
        // 说明：构建哈夫曼树后需要编码，直接写死在构造方法中，防止遗忘
        // 注意大括号，都要编码，不能放在else里面
        encode();
    }

    // 根据关系表恢复出对应的哈夫曼树，用于接受压缩文件后的初始化
    public HuffmanTree(Map<Character, StringBuilder> relation) {
        root = new HufNode('+');
        for (Character c : relation.keySet()) {
            StringBuilder path = relation.get(c);
            HufNode iter = root;
            for (int i = 0; i < path.length(); i++) {
                if (path.charAt(i) == '0') {
                    if (iter.left == null) iter.left = new HufNode('+');
                    iter = iter.left;
                } else {
                    if (iter.right == null) iter.right = new HufNode('+');
                    iter = iter.right;
                }
            }
            iter.value = c;
        }
    }

    // 编码哈夫曼树中的结点
    public Map<Character, StringBuilder> encode() {
        encode(root, new StringBuilder(""));
        return relation;
    }

    // 如果node为空，直接返回上一级，这样返回的时候字符串直接更新，不用判断
    // 添加到哈希表的时候注意要新建一个字符串，否则使用递归中的字符串，动态变化
    private void encode(HufNode node, StringBuilder res) {
        if (node == null) return;
        if (node.left == null && node.right == null) relation.put(node.value, new StringBuilder(res));
        encode(node.left, res.append('0'));
        res.delete(res.length() - 1, res.length());
        encode(node.right, res.append('1'));
        res.delete(res.length() - 1, res.length());
    }

    // 先编码原字符串，然后7位转asc码存储，最后记录最后一个字符用到的位数
    // 为了跟之前的普通字符区分开，使用非asc码字符李进行区分，至于为什么，哈哈哈
    public StringBuilder transfer() {
        StringBuilder temp = new StringBuilder("");
        for (int i = 0; i < origin.length; i++) {
            temp.append(relation.get(origin[i]));
        }
        StringBuilder res = new StringBuilder("");
        int count = 0, iter = 0, times = 1;
        // 确保有7位
        while (iter + 6 < temp.length()) {
            count = 0; times = 1;
            for (int i = 0; i < 7; i++) {
                count += times * (temp.charAt(iter) - '0');
                iter++; times *= 2;
            }
            res.append((char) count);
        }
        // 最后一位可能不足7，单独写
        count = 0; times = 1; int spare = temp.length() - iter;
        // 注意：这里正向遍历，但是times从1递增到64，所以相当于反过来的字符串转二进制
        // 因此这里times从1开始
        while (iter < temp.length()) {
            count += times * (temp.charAt(iter) - '0');
            iter++; times *= 2;
        }
        res.append((char) count);
        res.append("李").append(spare);
        return res;
    }

    // 将转换后的字符串恢复成原始字符串
    public StringBuilder decode(String transfer) {
        char[] arr = transfer.toCharArray();
        StringBuilder temp = new StringBuilder("");
        int iter = 0;

        // 处理分割线之前的字符，转成二进制
        // 之前转换的时候相当于字符串反过来，所以这里要颠倒
        while (arr[iter] != '李') {
            temp.append(Utils.char2Bin(arr[iter]).reverse());
            iter++;
        }

        // 删除多余字符
        iter++; temp.delete(temp.length() - 7 + transfer.charAt(iter) - '0',
                temp.length());

        // 遍历哈夫曼树恢复字符串
        arr = temp.toString().toCharArray();
        StringBuilder res = new StringBuilder("");
        iter = 0;
        while (iter < arr.length) {
            HufNode node = root;
            while (!(node.left == null && node.right == null)) {
                if (arr[iter] == '0') node = node.left;
                else node = node.right;
                iter++;
            }
            // 到根结点后取对应数据
            res.append(node.value);
        }
        return res;
    }

}
